…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion jasl#3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in jasl#1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in jasl#1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in jasl#1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion jasl#3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in jasl#1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in jasl#1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in jasl#1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>

…date loop

Redesigned suggestion #3 from PR vllm-project#41834 comment 4450901180. The first
attempt (e34daef, reverted; later 72a5ff2, also reverted) tried to
truncate ``topk_indices.shape[1]`` in Python so the captured launches
iterated a narrower combined slice; that approach broke under cudagraph
replay (shape baked at capture) and *also* mis-bounded — the combine
kernel writes each token's combined buffer as ``[topk_len_t |
swa_len_t | -1 padding]`` with SWA *immediately* following the
per-token topk, so a fixed ``effective_topk`` cap cuts off the SWA
portion (GSM8K dropped 25 pp on the prior attempt).

The kernel already loads the per-token combined length
(``valid_len = tl.load(lens_ptr + token_idx)`` for the four ``lens``-
gated kernels, ``gather_len`` for the two paged kernels). The existing
``is_valid`` guard only short-circuits the *heavy* work past that
length; the outer ``for candidate_idx in range(0, num_candidates)``
still pays one ``tl.load`` + branch per iter on the dead tail.
Capping the loop at
``min(num_candidates, valid_len - candidate_offset)`` (clamped to 0)
removes those wasted iterations while preserving the existing
``is_valid`` semantics: the iterations we now skip are exactly those
the existing guard already discarded.

Applied to six accumulate kernels in ``sparse_mla_kernels.py``:

- ``_accumulate_gathered_attention_chunk_kernel``
- ``_accumulate_indexed_attention_chunk_kernel``                 [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_kernel``      [autotuned in #1]
- ``_accumulate_fp8ds_global_slots_attention_chunk_multihead_kernel``  [decode]
- ``_accumulate_fp8ds_paged_attention_chunk_kernel``             [autotuned in #1]
- ``_accumulate_fp8ds_paged_attention_chunk_multihead_kernel``   [decode]

CUDA-graph safety: ``lens_ptr`` / ``gather_lens_ptr`` are stable
addresses; their values are refreshed per call by the metadata builder
(outside the captured forward) and by ``combine_topk_swa_indices``
(inside the forward but writing only into the persistent buffers the
accumulate kernels read from). The kernel inner-loop bound is a
runtime-loaded scalar — Triton compiles a dynamic loop and the captured
launch picks up the current value on each replay.

Savings scale with ``combined_topk_buffer_width - actual valid length``
(i.e. mostly visible at long ``max_model_len`` with shorter actual
contexts). At our test shape (``max_model_len=131072``, ISL=2048) the
saved iterations come mostly from the decode multihead path; expected
to be neutral / no-regression at short ``max_model_len`` where the
bound equals ``num_candidates``.

Signed-off-by: jasl <jasl9187@hotmail.com>